import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.impute import SimpleImputer
import xgboost as xgb
from sklearn.metrics import mean_squared_error, r2_score
import logging


class OlympicMedalPredictor:
    def __init__(self):
        # 初始化模型和数据处理组件
        self.athlete_model = None
        self.country_model = None
        self.sport_model = None
        self.scaler = StandardScaler()

        # 添加空值处理器
        self.imputer = SimpleImputer(strategy='mean')

        # 设置日志
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)

    def load_data(self):
        """加载所有预处理后的数据"""
        try:
            # 修改文件路径，从 features 文件夹读取
            data_path = "features"
            self.athlete_data = pd.read_csv(f'{data_path}/athlete_features.csv')
            self.country_data = pd.read_csv(f'{data_path}/country_features.csv')
            self.sport_data = pd.read_csv(f'{data_path}/sport_features.csv')
            self.time_series_data = pd.read_csv(f'{data_path}/time_series_features.csv')

            self.logger.info("数据加载成功")
            return True
        except Exception as e:
            self.logger.error(f"数据加载失败: {str(e)}")
            return False

    def prepare_athlete_features(self, data):
        """准备运动员特征"""
        features = pd.DataFrame()

        # 基础特征
        features['participation_count'] = data['participation_count']
        features['experience_years'] = data['last_year'] - data['first_year']
        features['avg_performance'] = data['avg_performance']
        features['best_performance'] = data['best_performance']
        features['total_medals'] = data['total_medals']

        # 计算近期表现趋势
        features['recent_performance'] = data['avg_performance'] * \
                                         (1 + np.log1p(data['participation_count']))

        # 计算经验权重
        features['experience_weight'] = np.log1p(features['experience_years'])

        return features

    def prepare_country_features(self, data):
        """准备国家特征"""
        features = pd.DataFrame()

        # 历史表现
        features['total_gold'] = data['total_gold']
        features['avg_gold'] = data['avg_gold']
        features['total_medals'] = data['total_medals']
        features['avg_medals'] = data['avg_medals']

        # 计算参与度
        features['participation_years'] = data['participation_years']
        features['host_advantage'] = data['times_as_host'].fillna(0)

        # 计算发展趋势
        features['medal_trend'] = (data['total_medals'] / data['participation_years']) * \
                                  np.log1p(data['participation_years'])

        return features

    def prepare_sport_features(self, data):
        """准备运动项目特征"""
        features = pd.DataFrame()

        # 竞争指标
        features['competition_index'] = data['competition_index']
        features['participating_countries'] = data['participating_countries']
        features['total_participants'] = data['total_participants']

        # 历史统计
        features['years_held'] = data['years_held']
        features['medal_winning_countries'] = data['medal_winning_countries']
        features['total_medals'] = data['total_medals']

        # 计算项目成熟度
        features['maturity_index'] = (data['years_held'] * data['participating_countries']) / \
                                     data['total_participants']

        return features

    def train_athlete_model(self, X, y):
        """训练运动员模型"""
        self.athlete_model = xgb.XGBRegressor(
            n_estimators=100,
            learning_rate=0.1,
            max_depth=5,
            random_state=42
        )
        self.athlete_model.fit(X, y)

    def train_country_model(self, X, y):
        """训练国家模型"""
        self.country_model = RandomForestRegressor(
            n_estimators=100,
            max_depth=10,
            random_state=42
        )
        self.country_model.fit(X, y)

    def train_sport_model(self, X, y):
        """训练运动项目模型"""
        self.sport_model = GradientBoostingRegressor(
            n_estimators=100,
            learning_rate=0.1,
            max_depth=4,
            random_state=42
        )
        self.sport_model.fit(X, y)

    def train_models(self):
        """训练所有模型"""
        # 准备特征
        athlete_features = self.prepare_athlete_features(self.athlete_data)
        country_features = self.prepare_country_features(self.country_data)
        sport_features = self.prepare_sport_features(self.sport_data)

        # 准备目标变量
        athlete_target = self.athlete_data['total_medals']
        country_target = self.country_data['total_medals']
        sport_target = self.sport_data['total_medals']

        # 处理特征中的空值
        self.logger.info("处理数据中的空值...")
        athlete_features = pd.DataFrame(self.imputer.fit_transform(athlete_features),
                                        columns=athlete_features.columns)
        country_features = pd.DataFrame(self.imputer.fit_transform(country_features),
                                        columns=country_features.columns)
        sport_features = pd.DataFrame(self.imputer.fit_transform(sport_features),
                                      columns=sport_features.columns)

        # 处理目标变量中的空值
        athlete_target = pd.Series(self.imputer.fit_transform(athlete_target.values.reshape(-1, 1)).ravel())
        country_target = pd.Series(self.imputer.fit_transform(country_target.values.reshape(-1, 1)).ravel())
        sport_target = pd.Series(self.imputer.fit_transform(sport_target.values.reshape(-1, 1)).ravel())

        # 划分训练集和测试集
        X_athlete_train, X_athlete_test, y_athlete_train, y_athlete_test = \
            train_test_split(athlete_features, athlete_target, test_size=0.2, random_state=42)

        X_country_train, X_country_test, y_country_train, y_country_test = \
            train_test_split(country_features, country_target, test_size=0.2, random_state=42)

        X_sport_train, X_sport_test, y_sport_train, y_sport_test = \
            train_test_split(sport_features, sport_target, test_size=0.2, random_state=42)

        # 标准化特征
        X_athlete_train = self.scaler.fit_transform(X_athlete_train)
        X_athlete_test = self.scaler.transform(X_athlete_test)

        X_country_train = self.scaler.fit_transform(X_country_train)
        X_country_test = self.scaler.transform(X_country_test)

        X_sport_train = self.scaler.fit_transform(X_sport_train)
        X_sport_test = self.scaler.transform(X_sport_test)

        # 训练模型
        self.train_athlete_model(X_athlete_train, y_athlete_train)
        self.train_country_model(X_country_train, y_country_train)
        self.train_sport_model(X_sport_train, y_sport_train)

        # 评估模型
        athlete_pred = self.athlete_model.predict(X_athlete_test)
        country_pred = self.country_model.predict(X_country_test)
        sport_pred = self.sport_model.predict(X_sport_test)

        # 打印评估结果
        self.logger.info("模型评估结果:")
        self.logger.info(f"运动员模型 R2 分数: {r2_score(y_athlete_test, athlete_pred):.4f}")
        self.logger.info(f"国家模型 R2 分数: {r2_score(y_country_test, country_pred):.4f}")
        self.logger.info(f"运动项目模型 R2 分数: {r2_score(y_sport_test, sport_pred):.4f}")

    def predict(self, athlete_features, country_features, sport_features):
        """组合预测"""
        # 处理空值
        athlete_features = pd.DataFrame(self.imputer.transform(athlete_features),
                                        columns=athlete_features.columns)
        country_features = pd.DataFrame(self.imputer.transform(country_features),
                                        columns=country_features.columns)
        sport_features = pd.DataFrame(self.imputer.transform(sport_features),
                                      columns=sport_features.columns)

        # 标准化输入特征
        athlete_features_scaled = self.scaler.transform(athlete_features)
        country_features_scaled = self.scaler.transform(country_features)
        sport_features_scaled = self.scaler.transform(sport_features)

        # 获取各模型预测
        athlete_pred = self.athlete_model.predict(athlete_features_scaled)
        country_pred = self.country_model.predict(country_features_scaled)
        sport_pred = self.sport_model.predict(sport_features_scaled)

        # 加权组合预测结果
        weights = [0.4, 0.4, 0.2]
        final_pred = (
                weights[0] * athlete_pred +
                weights[1] * country_pred +
                weights[2] * sport_pred
        )

        return final_pred

    def save_models(self, path):
        """保存训练好的模型"""
        import joblib

        try:
            joblib.dump(self.athlete_model, f"{path}/athlete_model.pkl")
            joblib.dump(self.country_model, f"{path}/country_model.pkl")
            joblib.dump(self.sport_model, f"{path}/sport_model.pkl")
            joblib.dump(self.scaler, f"{path}/scaler.pkl")
            joblib.dump(self.imputer, f"{path}/imputer.pkl")  # 保存空值处理器

            self.logger.info("模型保存成功")
            return True
        except Exception as e:
            self.logger.error(f"模型保存失败: {str(e)}")
            return False

    def load_models(self, path):
        """加载保存的模型"""
        import joblib

        try:
            self.athlete_model = joblib.load(f"{path}/athlete_model.pkl")
            self.country_model = joblib.load(f"{path}/country_model.pkl")
            self.sport_model = joblib.load(f"{path}/sport_model.pkl")
            self.scaler = joblib.load(f"{path}/scaler.pkl")
            self.imputer = joblib.load(f"{path}/imputer.pkl")  # 加载空值处理器

            self.logger.info("模型加载成功")
            return True
        except Exception as e:
            self.logger.error(f"模型加载失败: {str(e)}")
            return False


def main():
    # 1. 初始化预测器
    predictor = OlympicMedalPredictor()

    # 2. 加载数据
    if not predictor.load_data():
        print("数据加载失败，请检查数据文件是否存在")
        return

    # 3. 训练模型
    predictor.train_models()

    # 4. 保存模型
    predictor.save_models("saved_models")


if __name__ == "__main__":
    main()
